# link: https://mp.weixin.qq.com/s/NTct2_AYhz4Z8q5MYtBQcA

import backtrader as bt
import pandas as pd
import datetime
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


codes = ['sh600276', 'sh600519', 'sh603288']

# 恒瑞医药
data1 = get_data(codes[0])
# 贵州茅台
data2 = get_data(codes[1])
# 海天味业
data3 = get_data(codes[2])

# %%

# 实例化策略
cerebro = bt.Cerebro()

st_date = datetime.datetime(2023, 4, 4)
ed_date = datetime.datetime(2023, 5, 11)

# 添加 600276.SH 的行情数据
datafeed1 = bt.feeds.PandasData(dataname=data1, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed1, name=codes[0])

# 添加 600519.SH 的行情数据
datafeed2 = bt.feeds.PandasData(dataname=data2, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed2, name=codes[1])

# 添加 603288.SH 的行情数据
datafeed3 = bt.feeds.PandasData(dataname=data3, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed3, name=codes[2])


# 第一章 DataFeed的数据结构
# 第1.1节：验证 data 的结构
class TestStrategy(bt.Strategy):
    def __init__(self):
        # 打印数据集和数据集对应的名称
        print("-------------self.datas-------------")
        print(self.datas)
        print("-------------self.data-------------")
        print(self.data._name, self.data)  # 返回第一个导入的数据表格，缩写形式
        print("-------------self.data0-------------")
        print(self.data0._name, self.data0)  # 返回第一个导入的数据表格，缩写形式
        print("-------------self.datas[0]-------------")
        print(self.datas[0]._name, self.datas[0])  # 返回第一个导入的数据表格，常规形式
        print("-------------self.datas[1]-------------")
        print(self.datas[1]._name, self.datas[1])  # 返回第二个导入的数据表格，常规形式
        print("-------------self.datas[-1]-------------")
        print(self.datas[-1]._name, self.datas[-1])  # 返回最后一个导入的数据表格
        print("-------------self.datas[-2]-------------")
        print(self.datas[-2]._name, self.datas[-2])  # 返回倒数第二个导入的数据表格


cerebro.addstrategy(TestStrategy)
result = cerebro.run()
